"""
任务 房价价格预测

"""
import pandas as pd
from my_utils import *
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt

DATA_FILE = ".\data\house_data.csv"

# 使用的特征列
"""
bedrooms:房间个数  bathrooms:卫生间个数 sqft_living:居住平米数 sqft_lot:总平米数
sqft_above: 上层平米数 sqft_basement: 底层平米数
"""
FEAT_COLS = ['bedrooms','bathrooms','sqft_living','sqft_lot','sqft_above','sqft_basement']


def plot_fitting_line(linear_reg_model, X, y, feat):
    """
    绘制线性回归线
    :param linear_reg_model:
    :param X:
    :param y:
    :param feat:
    :return:
    """
    w = linear_reg_model.coef_
    b = linear_reg_model.intercept_
    plt.figure(figsize=(20,8),dpi=100)
    # 散点图 样本点
    plt.scatter(X, y,alpha=0.5)

    #直线
    plt.plot(X, w * X + b, c='red')
    plt.title(feat)
    plt.show()


def main():
    house_data = pd.read_csv(DATA_FILE, usecols=FEAT_COLS + ['price'])
    for feat in FEAT_COLS:
        X = house_data[feat].values.reshape(-1,1)
        y = house_data['price'].values
        X_train, X_test, y_train, y_test = train_test_split(X,y,test_size=1/3,random_state=10)
        linear_reg_model = LinearRegression()
        linear_reg_model.fit(X_train,y_train)
        r2_score = linear_reg_model.score(X_test,y_test)
        print("特征： {},R2值： {}".format(feat,r2_score))
        # 绘制拟合直线
        plot_fitting_line(linear_reg_model,X_train,y_train,feat)

if __name__ =='__main__':
    main()